# -*- coding: utf-8 -*-
# @time: 2024/6/5 9:47
# @file: index_backtest
# @author: tyshixi08

from get_data.origin_data import *

# 行列对齐(打印调试用)
pd.set_option('display.unicode.ambiguous_as_wide', True)
pd.set_option('display.unicode.east_asian_width', True)
pd.set_option('display.width', None)
# pd.set_option('display.max_rows', None)
pd.set_option('display.max_columns', None)

# 指标分位数回滚
def index_quantile(df_basic, window1, window2, sell_quantile, buy_quantile, flag=2):
    '''
    :param df_basic: 因子原始数据
    :param window1: 卖出滚动窗口
    :param window2: 买入滚动窗口
    :param sell_quantile: 卖出分位数
    :param buy_quantile: 买入分位数
    :param flag: 买入判断，1为自开始时滚动，2为按买入滚动窗口滚动
    :return: dataframe数据
    '''

    df = df_basic.copy()

    df = df.sort_values('date')
    df.set_index('date', inplace=True)

    if flag == 1:
        df[f'sell_quantile'] = df.iloc[:, 0].rolling(window = window1).quantile(sell_quantile)
        df[f'buy_quantile'] = df.iloc[:, 0].rolling(window=df.iloc[:, 0].shape[0], min_periods=1).quantile(buy_quantile)

    elif flag == 2:
        df[f'sell_quantile'] = df.iloc[:, 0].rolling(window = window1).quantile(sell_quantile)
        df[f'buy_quantile'] = df.iloc[:, 0].rolling(window = window2).quantile(buy_quantile)

    df = df.reset_index().dropna().reset_index(drop=True)
    return df

# 回测
def backtest(df, window1, window2, sell_quantile, buy_quantile, flag):
    '''
    :param df: 因子原始数据
    :param window1: 卖出滚动窗口
    :param window2: 买入滚动窗口
    :param sell_quantile: 卖出分位数
    :param buy_quantile: 买入分位数
    :param flag: 买入判断，1为自开始时滚动，2为按买入滚动窗口滚动
    :return: dataframe数据
    '''

    df_basic = df.copy()

    df_quantile = index_quantile(df_basic, window1, window2, sell_quantile, buy_quantile, flag)
    df_quantile.columns = ['交易日期', '使用指标', f'卖出{int(sell_quantile * 100)}%分位数', f'买入{int(buy_quantile * 100)}%分位数']
    df_basic_index = basic_index()
    df_basic_index['交易日期'] = pd.to_datetime(df_basic_index['交易日期']).dt.strftime('%Y-%m-%d')
    df = pd.merge(df_quantile, df_basic_index, how='outer', on='交易日期')
    df = df[['交易日期', '使用指标', '指数净值', f'卖出{int(sell_quantile * 100)}%分位数', f'买入{int(buy_quantile * 100)}%分位数']].dropna().sort_values('交易日期').reset_index(drop=True)

    # 获取持仓信号
    df['卖出信号'] = np.nan
    df['买入信号'] = np.nan

    # 按照输入的分位数滚动
    if flag == 2:
        df['持仓信号'] = np.nan
        df.loc[0, '持仓信号'] = 0
        for i in range(1, len(df)):
            if df.loc[i, '使用指标'] < df.loc[i, f'买入{int(buy_quantile * 100)}%分位数']:
                df.loc[i, '持仓信号'] = 1
            elif df.loc[i, '使用指标'] > df.loc[i, f'卖出{int(sell_quantile * 100)}%分位数']:
                df.loc[i, '持仓信号'] = 0
            else:
                df.loc[i, '持仓信号'] = df.loc[i - 1, '持仓信号']

    if flag == 1:
        for i in range(1, len(df)):
            if df.loc[i, '使用指标'] < df.loc[i, f'买入{int(buy_quantile * 100)}%分位数']:
                df.loc[i, '买入信号'] = 1
            if df.loc[i, '使用指标'] > df.loc[i, f'卖出{int(sell_quantile * 100)}%分位数']:
                df.loc[i, '卖出信号'] = 0

        df['持仓信号'] = np.nan
        df.loc[0, '持仓信号'] = 0

        for i in range(1,len(df)):
            if df.loc[i-1, '持仓信号'] == 0 and df.loc[i, '买入信号'] == 1:
                df.loc[i, '持仓信号'] = 1
            elif df.loc[i-1, '持仓信号'] == 1 and df.loc[i, '卖出信号'] == 0:
                df.loc[i, '持仓信号'] = 0
            else:
                df.loc[i, '持仓信号'] = df.loc[i - 1, '持仓信号']

    df.loc[0, '持仓状态'] = 0
    df['持仓状态'] = df['持仓信号'].shift(1)

    df = df[1:]

    df['指数净值'] = df['指数净值'] / df['指数净值'].iloc[0]

    df = df.reset_index(drop=True)

    df['日收益率'] = 0
    for i in range(1, len(df)):
        if df.loc[i, '持仓状态'] == 0:
            df.loc[i, '日收益率'] = 0
        else:
            df.loc[i, '日收益率'] = df.loc[i, '指数净值'] / df.loc[i - 1, '指数净值'] - 1

    df.loc[0, '累积净值'] = 1
    for i in range(1, len(df)):
        df.loc[i, '累积净值'] = df.loc[i - 1, '累积净值'] * (df.loc[i, '日收益率'] + 1)

    df = df[['交易日期', '指数净值', '使用指标', f'卖出{int(sell_quantile * 100)}%分位数', f'买入{int(buy_quantile * 100)}%分位数', '持仓信号', '持仓状态', '日收益率', '累积净值']]
    return df
